{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "header-cell",
   "metadata": {},
   "source": [
    "# 😄 原创IP生成器 - Qwen Image Edit\n",
    "\n",
    "基于 Qwen-Image-Edit 模型的原创IP生成工具"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install-deps",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 安装必要的依赖\n",
    "!pip install gradio diffusers modelscope accelerate transformers  openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f047979",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 检查版本信息\n",
    "from importlib.metadata import version\n",
    "print(\"diffusers version:\", version(\"diffusers\"))\n",
    "print(\"torch version:\", version(\"torch\"))\n",
    "print(\"transformers version:\", version(\"transformers\"))\n",
    "print(\"PIL version:\", version(\"Pillow\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "import torch\n",
    "from PIL import Image\n",
    "import os\n",
    "from diffusers import QwenImageEditPipeline\n",
    "from modelscope import snapshot_download\n",
    "import numpy as np\n",
    "import datetime\n",
    "import glob\n",
    "from typing import List, Optional\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "# 添加工具模块路径\n",
    "sys.path.append('./src/examples/tools')\n",
    "\n",
    "# 从 .env 文件加载环境变量\n",
    "def load_env_file():\n",
    "    env_file = Path('.env')\n",
    "    if env_file.exists():\n",
    "        with open(env_file, 'r', encoding='utf-8') as f:\n",
    "            for line in f:\n",
    "                line = line.strip()\n",
    "                if line and not line.startswith('#') and '=' in line:\n",
    "                    key, value = line.split('=', 1)\n",
    "                    key = key.strip()\n",
    "                    value = value.strip().strip('\"').strip(\"'\")\n",
    "                    os.environ[key] = value\n",
    "        print(\"✅ 已从 .env 文件加载环境变量\")\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "# 加载环境变量\n",
    "env_loaded = load_env_file()\n",
    "\n",
    "# 检查API密钥是否设置\n",
    "if 'OPENAI_API_KEY' not in os.environ:\n",
    "    print(\"⚠️ 未找到 OPENAI_API_KEY\")\n",
    "    if not env_loaded:\n",
    "        print(\"💡 请在项目根目录创建 .env 文件，内容如下:\")\n",
    "        print(\"   OPENAI_API_KEY=your-api-key-here\")\n",
    "    else:\n",
    "        print(\"💡 请在 .env 文件中添加 OPENAI_API_KEY=your-api-key-here\")\n",
    "    print(\"   或者设置系统环境变量 OPENAI_API_KEY\")\n",
    "else:\n",
    "    print(\"✅ OPENAI_API_KEY 已设置\")\n",
    "\n",
    "# 导入专用提示词生成工具\n",
    "try:\n",
    "    from ip_creation_prompt_utils import (\n",
    "        generate_creative_prompt, \n",
    "        get_subject_suggestions, \n",
    "        get_action_suggestions,\n",
    "        validate_inputs\n",
    "    )\n",
    "    print(\"✅ 成功导入专用提示词生成工具\")\n",
    "    IP_CREATION_AVAILABLE = True\n",
    "except ImportError as e:\n",
    "    print(f\"⚠️ 专用提示词生成工具导入失败: {e}\")\n",
    "    print(\"💡 请确保已设置 OPENAI_API_KEY 环境变量\")\n",
    "    IP_CREATION_AVAILABLE = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup-device",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置设备和数据类型\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "    torch_dtype = torch.bfloat16\n",
    "    print(f\"✅ 使用 GPU: {torch.cuda.get_device_name()}\")\n",
    "    print(f\"💾 GPU 内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB\")\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "    torch_dtype = torch.float32\n",
    "    print(\"⚠️ 使用 CPU（建议使用GPU以获得更好性能）\")\n",
    "\n",
    "print(f\"🔧 设备: {device}, 数据类型: {torch_dtype}\")\n",
    "\n",
    "# 创建输出目录\n",
    "output_dir = \"./ip_outputs\"\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "print(f\"📁 输出目录: {output_dir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "download-model",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下载并加载 Qwen-Image-Edit 模型\n",
    "model_id = \"Qwen/Qwen-Image-Edit\"\n",
    "local_dir = './models/Qwen-Image-Edit'\n",
    "\n",
    "# 检查模型是否已存在\n",
    "if not os.path.exists(local_dir):\n",
    "    print(f\"📥 开始下载 {model_id} 模型...\")\n",
    "    os.makedirs(os.path.dirname(local_dir), exist_ok=True)\n",
    "    snapshot_download(model_id, local_dir=local_dir)\n",
    "    print(f\"✅ 模型下载完成: {local_dir}\")\n",
    "else:\n",
    "    print(f\"✅ 模型已存在: {local_dir}\")\n",
    "\n",
    "# 加载管道\n",
    "print(\"🔄 正在加载图像编辑管道...\")\n",
    "pipeline = QwenImageEditPipeline.from_pretrained(\n",
    "    local_dir, \n",
    "    torch_dtype=torch_dtype,\n",
    "    use_safetensors=True,\n",
    "    device_map=\"balanced\"\n",
    ")\n",
    "\n",
    "# 验证它是实例对象\n",
    "print(type(pipeline))  # <class 'diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline'>\n",
    "print(isinstance(pipeline, QwenImageEditPipeline))  # True\n",
    "print(callable(pipeline))  # True，因为有 __call__ 方法\n",
    "\n",
    "# 可以访问实例属性和方法\n",
    "print(pipeline.vae)  # 访问 VAE 组件\n",
    "print(pipeline.transformer)  # 访问 Transformer 组件\n",
    "print(\"✅ 模型加载完成\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "emoji-functions",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_ip(image, prompt, seed=42, steps=50, cfg_scale=4.0):\n",
    "    \"\"\"\n",
    "    生成原创IP\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # 处理输入图像\n",
    "        if isinstance(image, np.ndarray):\n",
    "            image = Image.fromarray(image)\n",
    "        \n",
    "        if image.mode != \"RGB\":\n",
    "            image = image.convert(\"RGB\")\n",
    "        \n",
    "        # 编辑参数\n",
    "        inputs = {\n",
    "            \"image\": image,\n",
    "            \"prompt\": prompt,\n",
    "            \"generator\": torch.manual_seed(seed),\n",
    "            \"true_cfg_scale\": cfg_scale,\n",
    "            \"negative_prompt\": \"blurry, low quality, distorted\",\n",
    "            \"num_inference_steps\": steps,\n",
    "        }\n",
    "        \n",
    "        # 执行编辑\n",
    "        with torch.inference_mode():\n",
    "            output = pipeline(**inputs)\n",
    "            edited_image = output.images[0]\n",
    "        \n",
    "        # 保存生成的图像\n",
    "        timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "        filename = f\"ip_{timestamp}_{seed}.png\"\n",
    "        filepath = os.path.join(output_dir, filename)\n",
    "        edited_image.save(filepath)\n",
    "        \n",
    "        return edited_image\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ 生成原创IP时出错: {str(e)}\")\n",
    "        return image\n",
    "\n",
    "def get_history_images():\n",
    "    \"\"\"\n",
    "    获取历史生成的图片\n",
    "    \"\"\"\n",
    "    image_files = glob.glob(os.path.join(output_dir, \"*.png\"))\n",
    "    image_files.sort(key=os.path.getmtime, reverse=True)  # 按修改时间倒序\n",
    "    return image_files[:12]  # 返回最新的12张图片\n",
    "\n",
    "def load_history_image(image_path):\n",
    "    \"\"\"\n",
    "    加载历史图片\n",
    "    \"\"\"\n",
    "    if image_path and os.path.exists(image_path):\n",
    "        return Image.open(image_path)\n",
    "    return None\n",
    "\n",
    "def generate_ip_creation_prompt(subject, simple_prompt):\n",
    "    \"\"\"\n",
    "    生成专用的IP创作提示词\n",
    "    \"\"\"\n",
    "    if not IP_CREATION_AVAILABLE:\n",
    "        return f\"Make the {subject} {simple_prompt}\"\n",
    "    \n",
    "    if not subject or not subject.strip():\n",
    "        return \"请先输入图片主体\"\n",
    "    \n",
    "    if not simple_prompt or not simple_prompt.strip():\n",
    "        return \"请先输入简单动作\"\n",
    "    \n",
    "    # 验证输入\n",
    "    is_valid, error_msg = validate_inputs(subject.strip(), simple_prompt.strip())\n",
    "    if not is_valid:\n",
    "        return error_msg\n",
    "    \n",
    "    try:\n",
    "        generated = generate_creative_prompt(subject.strip(), simple_prompt.strip())\n",
    "        return generated\n",
    "    except Exception as e:\n",
    "        print(f\"⚠️ 专用提示词生成失败: {e}\")\n",
    "        return f\"Make the {subject} {simple_prompt}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "gradio-interface",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建原创IP生成器界面\n",
    "def create_ip_generator():\n",
    "    with gr.Blocks(title=\"原创IP生成器\", theme=gr.themes.Soft() ) as demo:\n",
    "        gr.Markdown(\n",
    "            \"\"\"\n",
    "            # 😄 原创IP生成器\n",
    "            \n",
    "            使用 Qwen-Image-Edit 模型，将普通图片转换为有趣的原创IP！\n",
    "            \n",
    "            **使用步骤：** 上传图片 → 输入图片主体 → 输入简单动作 → 生成专用提示词 → 生成原创IP\n",
    "            \n",
    "            **新功能：** 专用提示词生成器 - 只需输入主体和简单动作，AI会自动生成详细的创意提示词！\n",
    "            \"\"\"\n",
    "        )\n",
    "        \n",
    "        with gr.Row():\n",
    "            # 左侧：输入区域\n",
    "            with gr.Column(scale=1):\n",
    "                # 图片输入\n",
    "                input_image = gr.Image(\n",
    "                    label=\"图片输入\",\n",
    "                    type=\"pil\",\n",
    "                    height=300\n",
    "                )\n",
    "                \n",
    "                # 主体询问区域\n",
    "                gr.Markdown(\"**图片主体**\")\n",
    "                subject_input = gr.Textbox(\n",
    "                    label=\"\",\n",
    "                    placeholder=\"请输入图片中的主体，例如：小熊、猫、人物等...\",\n",
    "                    lines=1\n",
    "                )\n",
    "                \n",
    "                # 主体建议按钮\n",
    "                with gr.Row():\n",
    "                    btn_bear = gr.Button(\"小熊\", size=\"sm\")\n",
    "                    btn_cat = gr.Button(\"猫\", size=\"sm\")\n",
    "                    btn_dog = gr.Button(\"狗\", size=\"sm\")\n",
    "                    btn_rabbit = gr.Button(\"兔子\", size=\"sm\")\n",
    "                    btn_person = gr.Button(\"人物\", size=\"sm\")\n",
    "                \n",
    "                # 提示词区域\n",
    "                gr.Markdown(\"**简单动作**\")\n",
    "                \n",
    "                # 快速动作按钮\n",
    "                with gr.Row():\n",
    "                    btn_paint = gr.Button(\"画画\", size=\"sm\")\n",
    "                    btn_guitar = gr.Button(\"弹吉他\", size=\"sm\")\n",
    "                    btn_astronaut = gr.Button(\"宇航员\", size=\"sm\")\n",
    "                    btn_magician = gr.Button(\"魔法师\", size=\"sm\")\n",
    "                \n",
    "                with gr.Row():\n",
    "                    btn_read = gr.Button(\"读书\", size=\"sm\")\n",
    "                    btn_cook = gr.Button(\"做饭\", size=\"sm\")\n",
    "                    btn_sleep = gr.Button(\"睡觉\", size=\"sm\")\n",
    "                    btn_dance = gr.Button(\"跳舞\", size=\"sm\")\n",
    "                \n",
    "                # 简单动作输入\n",
    "                simple_prompt_input = gr.Textbox(\n",
    "                    label=\"\",\n",
    "                    placeholder=\"输入简单动作，例如：画画、弹吉他、宇航员...\",\n",
    "                    lines=1\n",
    "                )\n",
    "                \n",
    "                # 专用提示词生成按钮\n",
    "                with gr.Row():\n",
    "                    generate_prompt_button = gr.Button(\n",
    "                        \"🎯 生成专用提示词\", \n",
    "                        variant=\"secondary\", \n",
    "                        size=\"sm\",\n",
    "                        visible=IP_CREATION_AVAILABLE\n",
    "                    )\n",
    "                    if not IP_CREATION_AVAILABLE:\n",
    "                        gr.Markdown(\"💡 *设置 OPENAI_API_KEY 环境变量以启用专用提示词生成功能*\", elem_id=\"ip-creation-tip\")\n",
    "                \n",
    "                # 生成的详细提示词显示\n",
    "                prompt_input = gr.Textbox(\n",
    "                    label=\"生成的详细提示词\",\n",
    "                    placeholder=\"这里将显示生成的详细英文提示词...\",\n",
    "                    lines=4,\n",
    "                    interactive=True\n",
    "                )\n",
    "                \n",
    "                \n",
    "                # 高级参数\n",
    "                gr.Markdown(\"**高级参数**\")\n",
    "                with gr.Accordion(\"参数设置\", open=False):\n",
    "                    seed_input = gr.Slider(0, 1000, value=42, step=1, label=\"随机种子\")\n",
    "                    steps_input = gr.Slider(10, 50, value=50, step=5, label=\"推理步数\")\n",
    "                    cfg_scale_input = gr.Slider(1.0, 8.0, value=4.0, step=0.5, label=\"CFG缩放\")\n",
    "                \n",
    "                generate_button = gr.Button(\"🎨 生成原创IP\", variant=\"primary\", size=\"lg\")\n",
    "            \n",
    "            # 右侧：输出区域\n",
    "            with gr.Column(scale=1):\n",
    "                # 图片输出\n",
    "                output_image = gr.Image(label=\"图片输出\", height=300)\n",
    "                \n",
    "                # 历史图片 - 占满剩余空间\n",
    "                gr.Markdown(\"**历史图片**\")\n",
    "                history_gallery = gr.Gallery(\n",
    "                        label=\"\",\n",
    "                        show_label=False,\n",
    "                        elem_id=\"gallery\",\n",
    "                        columns=3,\n",
    "                        rows=4,\n",
    "                        height=500,\n",
    "                        object_fit=\"contain\",\n",
    "                        container=True,\n",
    "                        allow_preview=True\n",
    "                    )\n",
    "        \n",
    "        # 处理函数\n",
    "        def process_generation(image, prompt, seed, steps, cfg_scale, progress=gr.Progress(track_tqdm=True)):\n",
    "            if image is None:\n",
    "                gr.Warning(\"请先上传图片！\")\n",
    "                return None, get_history_images()\n",
    "            if not prompt.strip():\n",
    "                gr.Warning(\"请输入提示词！\")\n",
    "                return None, get_history_images()\n",
    "            \n",
    "            try:\n",
    "                # 生成原创IP\n",
    "                result_image = generate_ip(image, prompt, seed, steps, cfg_scale)\n",
    "                # 更新历史图片\n",
    "                history_images = get_history_images()\n",
    "                return result_image, history_images\n",
    "                \n",
    "            except Exception as e:\n",
    "                gr.Error(f\"生成失败: {str(e)}\")\n",
    "                return None, get_history_images()\n",
    "        \n",
    "        def update_history():\n",
    "            return get_history_images()\n",
    "        \n",
    "        def process_generate_ip_prompt(subject, simple_prompt):\n",
    "            \"\"\"\n",
    "            处理专用提示词生成\n",
    "            \"\"\"\n",
    "            if not subject or not subject.strip():\n",
    "                gr.Warning(\"请先输入图片主体！\")\n",
    "                return \"\"\n",
    "            \n",
    "            if not simple_prompt or not simple_prompt.strip():\n",
    "                gr.Warning(\"请先输入简单动作！\")\n",
    "                return \"\"\n",
    "            \n",
    "            try:\n",
    "                generated_prompt = generate_ip_creation_prompt(subject, simple_prompt)\n",
    "                gr.Info(\"✅ 专用提示词生成完成！\")\n",
    "                return generated_prompt\n",
    "            except Exception as e:\n",
    "                gr.Error(f\"专用提示词生成失败: {str(e)}\")\n",
    "                return f\"Make the {subject} {simple_prompt}\"\n",
    "        \n",
    "        # 绑定事件\n",
    "        generate_button.click(\n",
    "            fn=process_generation,\n",
    "            inputs=[input_image, prompt_input, seed_input, steps_input, cfg_scale_input],\n",
    "            outputs=[output_image, history_gallery]\n",
    "        )\n",
    "        \n",
    "        # 主体建议按钮绑定\n",
    "        subject_suggestions = {\n",
    "            btn_bear: \"小熊\",\n",
    "            btn_cat: \"猫\",\n",
    "            btn_dog: \"狗\",\n",
    "            btn_rabbit: \"兔子\",\n",
    "            btn_person: \"人物\"\n",
    "        }\n",
    "        \n",
    "        for btn, subject in subject_suggestions.items():\n",
    "            btn.click(fn=lambda x=subject: x, outputs=[subject_input])\n",
    "        \n",
    "        # 快速动作按钮绑定\n",
    "        action_suggestions = {\n",
    "            btn_paint: \"画画\",\n",
    "            btn_guitar: \"弹吉他\",\n",
    "            btn_astronaut: \"宇航员\",\n",
    "            btn_magician: \"魔法师\",\n",
    "            btn_read: \"读书\",\n",
    "            btn_cook: \"做饭\",\n",
    "            btn_sleep: \"睡觉\",\n",
    "            btn_dance: \"跳舞\"\n",
    "        }\n",
    "        \n",
    "        for btn, action in action_suggestions.items():\n",
    "            btn.click(fn=lambda x=action: x, outputs=[simple_prompt_input])\n",
    "        \n",
    "        # 专用提示词生成按钮绑定\n",
    "        if IP_CREATION_AVAILABLE:\n",
    "            generate_prompt_button.click(\n",
    "                fn=process_generate_ip_prompt,\n",
    "                inputs=[subject_input, simple_prompt_input],\n",
    "                outputs=[prompt_input]\n",
    "            )\n",
    "        \n",
    "        \n",
    "        # 页面加载时更新历史图片\n",
    "        demo.load(fn=update_history, outputs=[history_gallery])\n",
    "    \n",
    "    return demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "launch-app",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"🚀 启动原创IP生成器...\")\n",
    "    \n",
    "demo = create_ip_generator()\n",
    "\n",
    "# 启动 Gradio 应用\n",
    "demo.launch(\n",
    "    share=False,          # 不创建公共链接\n",
    "    server_name=\"0.0.0.0\",  # 允许外部访问\n",
    "    server_port=6006,    # 端口号\n",
    "    show_error=True,     # 显示详细错误信息\n",
    "    debug=True          # 调试模式\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
